#ifndef MYAMG_CSRPRECONDAMG_CUH
#define MYAMG_CSRPRECONDAMG_CUH
#include <CSRPrecond.cuh>
#include <JacobiSmoother.cuh>
#include <RugeStuben.cuh>
#include <utils/cudaProfiler.cuh>

namespace myamg {
    template<typename Integer, typename Real, typename Allocator, typename Smoother = JacobiSmoother<Integer, Real,
        Allocator> >
    class CSRPrecondAMG : public CSRPrecond<Integer, Real, Allocator, CSRPrecondAMG<Integer, Real, Allocator,
                Smoother> > {
    public:
        using matrixType = typename CSRPrecond<Integer, Real, Allocator, CSRPrecondAMG>::matrixType;
        using vectorType = typename CSRPrecond<Integer, Real, Allocator, CSRPrecondAMG>::vectorType;
        using integerType = typename CSRPrecond<Integer, Real, Allocator, CSRPrecondAMG>::integerType;
        using realType = typename CSRPrecond<Integer, Real, Allocator, CSRPrecondAMG>::realType;
        using allocatorType = typename CSRPrecond<Integer, Real, Allocator, CSRPrecondAMG>::allocatorType;
        using smootherType = Smoother;

        explicit CSRPrecondAMG(
            const int maxLevels_ = 20,
            const int coarseSize_ = 8,
            const double strengthThreshold_ = 0.25,
            const int aggressiveLevels_ = 0,
            const int coarsenType_ = 0,
            const int interpType_ = 0,
            const int interpMinElements_ = 4,
            const int interpMaxElements_ = 4,
            const double truncationFactor_ = 0.1,
            const double sparsificationThreshold_ = 0.01,
            const int cycleType_ = 0,
            const int preSweeps_ = 1,
            const int postSweeps_ = 1,
            const int coarseSweeps_ = 2,
            const double jacobiFactor_ = 0.75,
            const int SORType_ = 2,
            const double SORFactor_ = 1.0,
            const int ILUMaxFillins_ = 0,
            const double ILUDropTolerance_ = 0.01,
            const double chebyshevEigenRatio_ = 0.3,
            const bool printStats_ = false
        ) : maxLevels(maxLevels_), coarseSize(coarseSize_),
            strengthThreshold(strengthThreshold_), aggressiveLevels(aggressiveLevels_),
            coarsenType(coarsenType_), interpType(interpType_),
            interpMinElements(interpMinElements_), interpMaxElements(interpMaxElements_),
            truncationFactor(truncationFactor_), sparsificationThreshold(sparsificationThreshold_),
            cycleType(cycleType_), preSweeps(preSweeps_), postSweeps(postSweeps_),
            coarseSweeps(coarseSweeps_), jacobiFactor(jacobiFactor_), SORType(SORType_),
            SORFactor(SORFactor_), ILUMaxFillins(ILUMaxFillins_),
            ILUDropTolerance(ILUDropTolerance_), chebyshevEigenRatio(chebyshevEigenRatio_),
            printStats(printStats_), nLevels(0) {
        }

        void setupImpl(const matrixType &A_, const bool reuse = false) {
            if (reuse) {
                A[0] = A_;
                smoother[0].setup(A[0], true);
                return;
            }
            freeImpl();

            A.resize(maxLevels);
            P.resize(maxLevels - 1);
            R.resize(maxLevels - 1);
            r.resize(maxLevels);
            e.resize(maxLevels - 1);
            g.resize(maxLevels - 1);
            smoother.reserve(maxLevels);

            // 最细层行数与非零元数目
            size_t originalSize[2];
            originalSize[0] = A_.numRows();
            originalSize[1] = A_.numNonzeros();

            // 统计数据
            Real gridComplexity = 0.0;
            Real operatorComplexity = 0.0;

            // 初始化最细层
            A[0] = A_;

            // 当前层
            int level = 0;

            if (printStats) {
                std::cout << "Multigrid Hierarchy:\n";
                std::cout << "Level\tRows\tNz\n";
                std::cout << "--------------------------------------------------\n";
            }

            while (true) {
                // 初始化当前层平滑器
                smoother.emplace_back(buildSmoother());
                smoother[level].setup(A[level], false);

                // 当前层尺寸
                size_t currentSize[2];
                currentSize[0] = A[level].numRows();
                currentSize[1] = A[level].numNonzeros();
                const auto n = currentSize[0];

                // 更新统计数据
                gridComplexity += static_cast<Real>(currentSize[0]) / static_cast<Real>(originalSize[0]);
                operatorComplexity += static_cast<Real>(currentSize[1]) / static_cast<Real>(originalSize[1]);

                if (printStats) {
                    std::cout << level << "\t" << currentSize[0] << "\t" << currentSize[1] << "\n";
                }

                // 初始化r
                r[level] = vectorType(n);

                // 检查是否需要继续构造更粗层
                if (level == maxLevels - 1 || n <= static_cast<size_t>(coarseSize)) {
                    nLevels = level + 1;
                    break;
                }

                Vector<bool, allocatorType> stg(currentSize[1]);
                Vector<Integer, allocatorType> cfmap(currentSize[0]);

                // 计算强连接
                RugeStuben<integerType, realType, allocatorType>::strength(A[level], stg, strengthThreshold);
                // 粗化，生成C/F映射
                const bool isAggressive = (level < aggressiveLevels);
                coarsening(A[level], stg, cfmap, isAggressive);
                // 构造插值算子P
                interpolation(A[level], stg, cfmap, P[level], isAggressive);
                // 构造限制算子R = P^T
                R[level] = transposeMatrix(P[level]);

                // RAP
                matrixType C{};
                RugeStuben<integerType, realType, allocatorType>::RAP(R[level], A[level], P[level], C);

                // 稀疏化
                sparsification(C, A[level + 1]);

                e[level] = vectorType(A[level + 1].numRows());
                g[level] = vectorType(A[level + 1].numRows());

                ++level;
            }
            if (printStats) {
                std::cout << "--------------------------------------------------\n";
                std::cout << "Number of Levels: " << nLevels << "\n";
                std::cout << "Grid Complexity: " << gridComplexity << "\n";
                std::cout << "Operator Complexity: " << operatorComplexity << "\n";
            }
        }

        void freeImpl() {
            A.clear();
            P.clear();
            R.clear();
            smoother.clear();
            r.clear();
            e.clear();
            g.clear();
            nLevels = 0;
        }

        [[nodiscard]] size_t inSizeImpl() const {
            return A[0].numRows();
        }

        [[nodiscard]] size_t outSizeImpl() const {
            return A[0].numRows();
        }

        void applyImpl(const vectorType &b, vectorType &x) {
            switch (cycleType) {
                case 0: // V-cycle
                    vectorClear(x);
                    vCycle(0, b, x, true);
                    break;
                case 1: // W-cycle
                    vectorClear(x);
                    wCycle(0, b, x, true);
                    break;
                case 2: // F-cycle
                    fCycle(0, b, x);
                    break;
                default:
                    throw std::invalid_argument("Invalid cycle type.");
            }
        }

        matrixType& getA(const size_t level) {
            return A[level];
        }

        const matrixType& getA(const size_t level) const {
            return A[level];
        }

        matrixType& getP(const size_t level) {
            return P[level];
        }

        const matrixType& getP(const size_t level) const {
            return P[level];
        }

        matrixType& getR(const size_t level) {
            return R[level];
        }

        const matrixType& getR(const size_t level) const {
            return R[level];
        }

        vectorType& getr(const size_t level) {
            return r[level];
        }

        const vectorType& getr(const size_t level) const {
            return r[level];
        }

        vectorType& gete(const size_t level) {
            return e[level];
        }

        const vectorType& gete(const size_t level) const {
            return e[level];
        }

        vectorType& getg(const size_t level) {
            return g[level];
        }

        const vectorType& getg(const size_t level) const {
            return g[level];
        }

        [[nodiscard]] size_t getLevels() const {
            return nLevels;
        }

        int maxLevels; //!< 最大网格层数（默认为20）
        int coarseSize; //!< 最粗网格的大小，当网格大小小于该值时，停止构造多重网格（默认为8）
        double strengthThreshold; //!< 强度阈值，用于确定强连接（默认为0.25）
        int aggressiveLevels; //!< 激进聚合的层数（默认为0）
        int coarsenType; //!< 粗化算法，0（默认）：HMPI，1：PMIS
        int interpType; //!< 插值算法，0（默认）：Long Range Interpolation，1：Smoothed Aggregation
        int interpMinElements; //!< 最小插值数目（默认为4）
        int interpMaxElements; //!< 最大插值数目（默认为4）
        double truncationFactor; //!< 插值截断因子（默认为0.1）
        double sparsificationThreshold; //!< 粗网格矩阵稀疏化阈值（默认为0.01）
        int cycleType; //!< 多重网格循环类型，0（默认）：V循环，1：W循环，2：F循环
        int preSweeps; //!< 预平滑步数（默认为1）
        int postSweeps; //!< 后平滑步数（默认为1）
        int coarseSweeps; //!< 最粗网格平滑步数（默认为2）
        double jacobiFactor; //!< Jacobi松弛因子（默认为0.75）
        int SORType; //!< SOR类型，0：前向GS，1：后向GS，2（默认）：对称GS
        double SORFactor; //!< SOR松弛因子（默认为1.0）
        int ILUMaxFillins; //!< ILU最大填充级别（默认为0）
        double ILUDropTolerance; //!< ILU调降容差（默认为0.01）
        double chebyshevEigenRatio; //!< Chebyshev平滑的特征值比率（默认为0.3）
        bool printStats; //!< 是否打印多重网格统计信息（默认为false）

    private:
        void vCycle(const int level, const vectorType &b, vectorType &x, bool x0Zero = false) {
            if (level == nLevels - 1) {
                coarseSmooth(level, b, x, x0Zero);
            } else {
                preSmoothRestriction(level, b, x, g[level], x0Zero);

                vectorClear(e[level]);

                vCycle(level + 1, g[level], e[level], true);

                prolongation(level, e[level], x);

                postSmooth(level, b, x);
            }
        }

        void wCycle(const int level, const vectorType &b, vectorType &x, bool x0Zero = false) {
            if (level == nLevels - 1) {
                coarseSmooth(level, b, x, x0Zero);
            } else {
                preSmoothRestriction(level, b, x, g[level], x0Zero);

                vectorClear(e[level]);

                wCycle(level + 1, g[level], e[level], true);
                wCycle(level + 1, g[level], e[level]);

                prolongation(level, e[level], x);

                postSmooth(level, b, x);
            }
        }

        void fCycle(const int level, const vectorType &b, vectorType &x) {
            if (level == nLevels - 1) {
                vectorClear(x);
                coarseSmooth(level, b, x, true);
            } else {
                matMulVec(R[level], b, g[level]); // g = R * b

                fCycle(level + 1, g[level], e[level]);

                matMulVec(P[level], e[level], x); // x = P * e

                vCycle(level, b, x);
            }
        }

        void coarseSmooth(const int level, const vectorType &b, vectorType &x, bool x0Zero = false) {
            smoother[level].pre(A[level], b, x, coarseSweeps, x0Zero);
        }

        void preSmoothRestriction(const int level, const vectorType &b, vectorType &x, vectorType &g,
                                  bool x0Zero = false) {
            smoother[level].preResid(A[level], b, x, r[level], preSweeps, x0Zero);
            matMulVec(R[level], r[level], g); // g = R * r
        }

        void postSmooth(const int level, const vectorType &b, vectorType &x) {
            smoother[level].post(A[level], b, x, postSweeps);
        }

        void prolongation(const int level, const vectorType &e, vectorType &x) {
            // x = x + P[level] * e
            const size_t n = x.size();
            if constexpr (AllocatorType<allocatorType>::value == MemoryType::Host) {
                prolongationHost(x.values(), e.values(), P[level].rowPtr(), P[level].colIndex(), P[level].rowPtr(), n);
            } else if constexpr (AllocatorType<allocatorType>::value == MemoryType::Cuda) {
                dim3 block(512);
                dim3 grid((n + block.x - 1) / block.x);
                prolongationDevice<<<grid, block>>>(x.values(), e.values(), P[level].rowPtr(), P[level].colIndex(),
                                                    P[level].values(), n);
            } else {
                throw std::invalid_argument("Unsupported memory type for prolongation.");
            }
        }

        smootherType buildSmoother() {
            if constexpr (std::is_same_v<smootherType, JacobiSmoother<integerType, realType, allocatorType> >) {
                return smootherType(jacobiFactor);
            } else {
                return smootherType();
            }
        }

        void coarsening(const matrixType &A, const Vector<bool, allocatorType> &stg, Vector<int, allocatorType> &cfmap,
                        const bool isAggressive) {
            switch (coarsenType) {
                case 0: // HMPI
                    if (isAggressive) {
                        RugeStuben<integerType, realType, allocatorType>::AggressiveHMISCoarsening(A, stg, cfmap);
                    } else {
                        RugeStuben<integerType, realType, allocatorType>::HMISCoarsening(A, stg, cfmap);
                    }
                    break;
                case 1: //  PMIS
                    if (isAggressive) {
                        RugeStuben<integerType, realType, allocatorType>::AggressivePMISCoarsening(A, stg, cfmap);
                    } else {
                        RugeStuben<integerType, realType, allocatorType>::PMISCoarsening(A, stg, cfmap);
                    }
                    break;
                default:
                    throw std::invalid_argument("Invalid coarsening type.");
            }
        }

        void interpolation(const matrixType &A, const Vector<bool, allocatorType> &stg,
                           Vector<int, allocatorType> &cfmap, matrixType &P, const bool isAggressive) {
            switch (interpType) {
                case 0:
                    if (isAggressive) {
                        RugeStuben<integerType, realType, allocatorType>::aggressiveLongRangeInterpolation(A, stg, cfmap, P,
                                                                                                         interpMinElements,
                                                                                                         interpMaxElements,
                                                                                                         truncationFactor);
                    } else {
                        RugeStuben<integerType, realType, allocatorType>::longRangeInterpolation(A, stg, cfmap, P,
                                                                                                 interpMinElements,
                                                                                                 interpMaxElements,
                                                                                                 truncationFactor);
                    }
                    break;
                case 1:
                    if (isAggressive) {
                        RugeStuben<integerType, realType, allocatorType>::aggressiveSmoothedAggregation(A, stg, cfmap, P);
                    } else {
                        RugeStuben<integerType, realType, allocatorType>::smoothedAggregation(A, stg, cfmap, P);
                    }
                    break;
                default:
                    throw std::invalid_argument("Invalid interpolation type.");
            }
        }

        void sparsification(const matrixType& A, matrixType& B) {
            RugeStuben<integerType, realType, allocatorType>::sparsification(A, B, sparsificationThreshold);
        }

        int nLevels; //!< 实际构造的网格层数
        std::vector<matrixType> A; //!< 多重网格各层的矩阵
        std::vector<matrixType> P; //!< 多重网格各层的插值算子
        std::vector<matrixType> R; //!< 多重网格各层的限制算子
        std::vector<smootherType> smoother; //!< 多重网格各层的平滑器

        std::vector<vectorType> r; //!< 各层的残差向量
        std::vector<vectorType> e; //!< 各层的误差向量
        std::vector<vectorType> g; //!< 各层的解向量
    };
}

#endif //MYAMG_CSRPRECONDAMG_CUH
